from mindquantum import Circuit
from mindquantum.core.gates import H, X
from mindquantum.core.circuit import UN
from mindquantum.simulator import Simulator
import numpy as np

def create_ghz(drawing=False):
    """ Generate a Circuit for a 3 qubit GHZ State and measure it. If `drawing` is True, return both the circuit object and the Matplotlib drawing of the circuit, otherwise return just the circuit object.
    """
    ghz = Circuit()
    ghz += H.on(0)
    ghz += X.on(1, 0)
    ghz += X.on(2, 0)
    ghz.measure_all()
    if drawing:
        fig = ghz.draw()
        return ghz, fig
    return ghz

def check(candidate):
    import math
    import matplotlib

    def check_circuit(circuit):
        assert circuit[-1].name == "measure"
        circuit_without_measure = circuit.copy()
        circuit_without_measure.measure_reset()
        sim = Simulator('mqvector', 3)
        sim.apply_circuit(circuit_without_measure)
        state = sim.get_qs()
        expected_state = np.zeros(8, dtype=complex)
        expected_state[0] = 1 / math.sqrt(2)
        expected_state[7] = 1 / math.sqrt(2)
        assert np.allclose(np.abs(state), np.abs(expected_state), atol=1e-6)

    circuit = candidate()
    check_circuit(circuit)

    circuit, drawing = candidate(drawing=True)
    check_circuit(circuit)
    assert isinstance(drawing, matplotlib.figure.Figure)

    print("OK")

if __name__ == '__main__':
    check(create_ghz)